package ffte.native

import breeze.math.{Complex}

import ffte.algorithm.FFTGen.{Method,genericFT,DFT,psbFFT,msbFFT,wpbFFT}
import ffte.algorithm.{circularOrder,FFTMisc}

/**
 * Native FFT trait
 *
 * This trait defines the interface for native software FFT implementations.
 * It provides a common interface for different FFT algorithm implementations
 * that operate on complex arrays with configurable scaling.
 *
 * Key Features:
 * - Configurable scaling factor for fixed-point simulation
 * - Pure software implementation using Breeze complex numbers
 * - Support for various FFT algorithms (DFT, PSB, MSB, WPB)
 * - Reference implementation for hardware validation
 *
 * @param S Scaling factor (right shift amount) for fixed-point simulation
 */
trait NativeFFT {
    /** Scaling factor for fixed-point arithmetic simulation */
    val S  : Int

    /**
     * Perform FFT transformation
     *
     * Applies the specific FFT algorithm implementation to the input
     * complex array and returns the transformed result.
     *
     * @param x Input complex array
     * @return FFT-transformed complex array with applied scaling
     */
    def transform(x:Array[Complex]) : Array[Complex]
}

/**
 * Native FFT method generator
 *
 * This class implements the Method trait for generating native software FFT
 * implementations. It dispatches to specific algorithm implementations based
 * on the FFT structure type.
 *
 * Supported Algorithms:
 * - DFT: Direct DFT implementation
 * - psbFFT: Prime Space Base FFT
 * - msbFFT: Mixed Space Base FFT
 * - wpbFFT: Winograd Prime-Based FFT
 *
 * @param shift Global scaling factor for fixed-point simulation
 */
class Native(shift:Int) extends Method {

    /**
     * Generate native FFT implementation
     *
     * Creates the appropriate native FFT implementation based on the
     * algorithm type and structure.
     *
     * @param entity Generic FFT structure with algorithm-specific parameters
     * @param prefix Instance name for debugging/identification
     * @return Native FFT implementation instance
     */
    def apply(entity: genericFT, prefix: String):NativeFFT = {
        entity match {
            case d : DFT            => new ndft(d.N,shift)
            case d : psbFFT[Native] => new psb(shift,d)
            case d : msbFFT[Native] => new msb(shift,d)
            case d : wpbFFT[Native] => new wpb(shift,d)
            case _                  => new ndft(2,shift)
        }
    }
}

/**
 * Native Direct DFT implementation
 *
 * This class provides a native software implementation of the Direct DFT
 * algorithm using the mathematical definition. It serves as a reference
 * implementation and fallback for small transform sizes.
 *
 * Mathematical Definition:
 * X[k] = Σ_{n=0}^{N-1} x[n] * W_N^{k*n}
 * where W_N = e^(-j*2π/N)
 *
 * Characteristics:
 * - O(N²) computational complexity
 * - Accurate reference implementation
 * - Configurable scaling for fixed-point simulation
 * - Used for validation and small transforms
 *
 * @param N Transform size
 * @param S Scaling factor (right shift amount)
 */
sealed class ndft(N:Int,val S:Int) extends NativeFFT {

    /**
     * Direct DFT computation
     *
     * Implements the DFT using the direct mathematical definition.
     * Uses twiddle factors generated by FFTMisc.wnk for complex
     * exponential calculations.
     *
     * @param x Input complex array
     * @return DFT result without scaling
     */
    def dft(x:Array[Complex]) = {
		val n = x.length
		(0 until n).map{i => x.zipWithIndex.map{ case (a,k) => a*FFTMisc.wnk(i*k,n)}.sum}.toArray
	}

    /**
     * Transform with scaling
     *
     * Applies the DFT computation and then scales the result by
     * the specified scaling factor to simulate fixed-point arithmetic.
     *
     * @param x Input complex array
     * @return Scaled DFT result
     * @throws AssertionError if input length doesn't match expected size N
     */
    def transform(x:Array[Complex]) = {
        assert(x.length==N,s"Wrong number of data")
        dft(x).map(_/(1<<S))
    }
}

/**
 * Native Prime Space Base (PSB) FFT implementation
 *
 * This class implements the PSB FFT algorithm in native software. PSB
 * decomposition uses the Chinese Remainder Theorem for optimal data
 * reordering when factors are coprime.
 *
 * Mathematical Foundation:
 * - Uses Chinese Remainder Theorem for index mapping
 * - Requires coprime factors for optimal performance
 * - Two-stage decomposition with data reordering
 *
 * Algorithm Structure:
 * 1. First stage: Apply FFT to first dimension groups
 * 2. Data reordering using CRT-based index mapping
 * 3. Second stage: Apply FFT to second dimension groups
 *
 * @param S Global scaling factor
 * @param p PSB FFT structure with decomposition parameters
 * @tparam T Method type for recursive FFT generation
 */
sealed class psb[T<:Method](val S:Int,val p:psbFFT[T]) extends NativeFFT {

    /** First dimension FFT block */
    val first  = p.ib

    /** Second dimension FFT block */
    val second = p.ob

    /** Distributed scaling for first stage */
    val first_shift = Math.round(S*Math.log(first.N)/Math.log(second.N*first.N)).toInt

    /** Distributed scaling for second stage */
    val second_shift = S - first_shift

    /** Native FFT generator for first stage */
    val fg = new Native(first_shift)

    /** Native FFT generator for second stage */
    val sg = new Native(second_shift)

    /** First stage FFT computation block */
    val f_block = fg(first,"f")

    /** Second stage FFT computation block */
    val s_block = sg(second,"s")

    /**
     * PSB FFT transformation
     *
     * Implements the two-stage PSB algorithm with CRT-based data reordering.
     * Distributes scaling across stages for optimal precision.
     *
     * @param x Input complex array
     * @return PSB FFT-transformed complex array
     */
    def transform(x:Array[Complex]): Array[Complex] = {
        val r = new Array[Complex](x.length)

        // First stage: process first dimension groups
        for(i <- 0 until second.N) {
            val fout = f_block.transform(x.slice(i*first.N,(i+1)*first.N))
            for(k <- 0 until first.N) {
                r(k*second.N+i) = fout(k)
            }
        }

        // Second stage: process second dimension groups
        for(i <- 0 until first.N) {
            val fout = s_block.transform(r.slice(i*second.N,(i+1)*second.N))
            for(k <- 0 until second.N) {
                r(i*second.N+k) = fout(k)
            }
        }
        r
    }
}

/**
 * Native Mixed Space Base (MSB) FFT implementation
 *
 * This class implements the MSB FFT algorithm in native software. MSB
 * decomposition uses mixed-radix factorization with rotation factors
 * for cases where factors are not coprime.
 *
 * Mathematical Foundation:
 * - Uses mixed-radix factorization for composite sizes
 * - Applies rotation factors (twiddle factors) between stages
 * - Handles non-coprime factors efficiently
 * - Two-stage decomposition with complex multiplications
 *
 * Algorithm Structure:
 * 1. First stage: Apply FFT to first dimension groups
 * 2. Apply rotation factors to align data for second stage
 * 3. Second stage: Apply FFT to second dimension groups
 *
 * Key Features:
 * - Rotation factors computed using load/store permutations
 * - Handles cases where PSB (coprime factorization) is not applicable
 * - Distributed scaling across stages for optimal precision
 * - More general than PSB but with additional multiplicative overhead
 *
 * @param S Global scaling factor
 * @param p MSB FFT structure with decomposition and rotation parameters
 * @tparam T Method type for recursive FFT generation
 */
sealed class msb[T<:Method](val S:Int,val p:msbFFT[T]) extends NativeFFT {

    /** First dimension FFT block */
    val first  = p.ib

    /** Second dimension FFT block */
    val second = p.ob

    /** Distributed scaling for first stage */
    val first_shift = Math.round(S*Math.log(first.N)/Math.log(second.N*first.N)).toInt

    /** Distributed scaling for second stage */
    val second_shift = S - first_shift

    /** Native FFT generator for first stage */
    val fg = new Native(first_shift)

    /** Native FFT generator for second stage */
    val sg = new Native(second_shift)

    /** First stage FFT computation block */
    val f_block = fg(first,"f")

    /** Second stage FFT computation block */
    val s_block = sg(second,"s")

    /**
     * Rotation factors for MSB algorithm.
     *
     * These twiddle factors align the data between the two FFT stages
     * when the factors are not coprime. Computed using the load and
     * store permutation tables from the decomposition structure.
     */
    val rotator = (for(i <- 0 until second.N;k <- 0 until first.N) yield
        FFTMisc.wnk(p.p.first.store_tab(k)*p.p.second.load_tab(i),p.N)).toArray

    /**
     * MSB FFT transformation
     *
     * Implements the two-stage MSB algorithm with rotation factors.
     * Applies distributed scaling and complex multiplications for
     * non-coprime factor handling.
     *
     * @param x Input complex array
     * @return MSB FFT-transformed complex array
     */
    def transform(x:Array[Complex]): Array[Complex] = {
        val r = new Array[Complex](x.length)

        // First stage: process first dimension groups
        for(i <- 0 until second.N) {
            val fout = f_block.transform(x.slice(i*first.N,(i+1)*first.N))
            for(k <- 0 until first.N) {
                // Apply rotation factor during data reordering
                r(k*second.N+i) = fout(k)*rotator(k+i*first.N)
            }
        }

        // Second stage: process second dimension groups
        for(i <- 0 until first.N) {
            val fout = s_block.transform(r.slice(i*second.N,(i+1)*second.N))
            for(k <- 0 until second.N) {
                r(i*second.N+k) = fout(k)
            }
        }
        r
    }
}

/**
 * Native Winograd Prime-Based (WPB) FFT implementation
 *
 * This class implements the WPB FFT algorithm in native software. WPB
 * algorithms use Winograd minimal multiplication techniques combined
 * with prime factorization for optimal performance.
 *
 * Mathematical Foundation:
 * - Winograd minimal multiplication algorithms
 * - Circular order permutations for optimal structure
 * - Prime factorization for decomposition
 * - Special handling of DC component separation
 *
 * Algorithm Structure:
 * 1. Separate DC component and scaling
 * 2. Apply circular order permutation
 * 3. Winograd processing with pre/post-processing
 * 4. Final combination and DC restoration
 *
 * @param S Global scaling factor
 * @param p WPB FFT structure with Winograd parameters
 */
sealed class wpb(val S:Int,p:wpbFFT[Native]) extends NativeFFT {

    /** Circular order permutation structure */
    val oc = circularOrder(p.N)

    /** Winograd coefficients for frequency correction */
    val wf = oc.wcoeff

    /** Distributed scaling for first stage */
    val first_shift  = S/2

    /** Distributed scaling for second stage */
    val second_shift = S-first_shift

    /** Native FFT generator for first stage */
    val fg = new Native(first_shift)

    /** Native FFT generator for second stage */
    val sg = new Native(second_shift)

    /** First stage Winograd processing block */
    val f_block = fg(p.ib,"f")

    /** Second stage Winograd processing block */
    val s_block = sg(p.ob,"s")

    /** Frequency correction coefficients */
    val coeff = (for(i <- 0 until p.N-1) yield wf(p.p.first.store_tab(i))).toArray

    /**
     * WPB FFT transformation
     *
     * Implements the Winograd-based FFT algorithm with special DC handling
     * and frequency correction. Uses circular ordering for optimal
     * Winograd algorithm structure.
     *
     * @param x Input complex array
     * @return WPB FFT-transformed complex array
     */
    def transform(x: Array[Complex]): Array[Complex] = {
        // Compute scaled sum (DC component)
        val sum = (Complex(0.0,0.0) /: x)(_+_) /(1<<S)

        // Extract and scale DC component
        val x0  = x(0) / (1<<S)

        // Apply Winograd processing to non-DC components
        val tf  = f_block.transform(x.slice(1,x.length))

        // Apply frequency correction coefficients
        val wtf = tf.zipWithIndex.map{ case(x,i) => x*coeff(i)}

        // Second stage processing
        val ts  = s_block.transform(wtf)

        // Reconstruct output with DC restoration
        val r   = new Array[Complex](x.length)
        r(0) = sum
        ts.zipWithIndex.foreach{ case(x,i) => r(i+1) = x+x0}
        r
    }
}